# 扩写海关官网的QA，以获得与目标对话相关性更高的chatlog
import os
import json
import numpy as np
from openai import OpenAI

file_path = os.path.abspath(__file__)
dir_path = os.path.dirname(file_path)




def expand_chatlog(QA_summary):
    question = QA_summary["留言内容"]
    answer = QA_summary["回复内容"]
    user_content = """You are a language generation expert who specializes in inferring and restoring real conversation transcripts based on conversation summaries. I will provide you with a summary of conversations summarized by a user and customs customer service, each summary in the form of a question and answer pair (QA).\nI would like you to generate content for the restored conversation scenario based entirely on the summaries I provide. You can augment it by adding greetings between the customer service and the user, splitting the QA pairs to form multiple conversations, and so on. The final dialog scene should be limited to between 15-25 sentences.\nNote: You cannot add additional information that may have an impact on the dialog summary.
Here is what the article summary will look like:"""
    user_content += f"'''question: {question}\nanswer: {answer}'''"

    client = OpenAI(
        api_key="sk-81bf128c90424cd5bec2c9a3c54ef309",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    messages = [
        {
            "role": "system",
            "content": "You're an excellent scene reconstructionist for dialog notes.",
        },
        {
            "role": "user",
            "content": user_content,
        },
        {
            "role": "assistant",
            "content": "",
        },
    ]
    response = client.chat.completions.create(
        model="qwen-max",  # 指定使用的模型名称
        messages=messages,  # 定义的消息列表
        max_tokens=1024,  # 限制最大令牌数为150
        response_format={
            "type": "text",  # 响应类型为纯文本
        }
    )
    ret = response.choices[0].message.content
    return ret




def main():
    with open(os.path.join(dir_path, "QA_test_data.json"), "r", encoding="utf-8") as f:
        QA_summaries = json.load(f)

    total_summaries = len(QA_summaries)
    np.random.seed(999)
    idx = np.random.randint(total_summaries)
    QA_summary = QA_summaries[idx]

    QA_chatlog = expand_chatlog(QA_summary)

    with open(os.path.join(dir_path, "../output/QA_expend_demo_01.txt"), "w", encoding="utf-8") as f:
        f.write(QA_summary["留言内容"])
        f.write("\n")
        f.write("\n")
        f.write("\n")
        f.write("\n")
        f.write("\n")
        f.write(QA_summary["回复内容"])
        f.write("\n")
        f.write("\n")
        f.write("\n")
        f.write("\n")
        f.write("\n")
        f.write("*"*100)
        f.write("\n")
        f.write(QA_chatlog)
    


if __name__ == "__main__":
    main()